
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
f
import matplotlib.pyplot as plt
import sys
sys.path.append("models/")
from custom_model import CNN5NET
from ghostnetv2_torch import ghostnetv2
from  ShuffleNetv2_torch import ShuffleNetV2
import time
import os
from math import cos, pi

from torch.utils.tensorboard import SummaryWriter
# %matplotlib inline
# **************************************************************

# **************************************************************

checkpoint = "output/take_photo/ghostnetv2_224_224_classify_2class_cocoadd_0128/2024_01_28_23_13_43/best.pth"
savefolder = os.path.dirname(checkpoint) + "/ghostnetv2_224_224_classify_2class_cocoadd_20240128.onnx"


# CPU 或者 GPU
device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')
# 初始化网络,加载预训练模型
# model = CNN5NET(num_classes=14, init_weights=False)
model = ghostnetv2(num_classes=2, 
                       width=1, 
                       dropout=0.2)
# model = ShuffleNetV2(num_classes=2)
state_dict = torch.load(checkpoint, map_location='cpu')

model.load_state_dict(state_dict['model'])


x = torch.randn(1, 3, 224, 224)
with torch.no_grad():
    torch.onnx.export(
        model,
        x,
        savefolder,
        opset_version=12,
        input_names=['input_0'],
        output_names=['output_0'],
        verbose = False
    )

